import json

from json_repair import repair_json
from loguru import logger

from utils.comm1 import time_count

def fix_json_string(input_string):
  s1 = input_string.replace('”', '"').replace("，", ",").replace("：", ":")  # 替换全角双引号和全角逗号

  import re
  pattern = r"(:\s?True)"  # 替换 : True  --》 : true
  s2 = re.sub(pattern, ": true", s1)

  pattern = r"(:\s?False)"  # 替换 : False  --》 : false
  s3 = re.sub(pattern, ": true", s2)
  return s3

def parse_chat_response(input_string):
  import re
  import json
  result = re.search(r"\{(.+)\}", input_string, re.DOTALL)
  if result:
    extracted_content = fix_json_string(result.group(1))
    fix_str1 = "{" + extracted_content + "}"
    try:
      json_data = json.loads(fix_str1)
      return json_data
    except Exception as e:
      logger.warning(f"json_data error , try to repair ... ")
      fix_str2 = repair_json(fix_str1)
      try:
        json_data = json.loads(fix_str2)
        return json_data
      except Exception as e:
        logger.warning(f"json_data error: {e}")
        logger.warning(f"无法解析成对象，返回的原始字符串： {fix_str1}")
        return fix_str1
  else:
    return None

def try_parse_json_datas(input_string):
  """
  尝试解析以下这种格式的json字符串：
  {"a":1}
  {"b":2}
  并返回 对象 List
  """
  logger.info("try_parse_json_datas ... ")
  import re
  spec_line_split = "\n\n_block_split_\n\n"
  # pattern = r".*?\}\s*?\{"
  pattern = r".*?\}\s*[,]*\s*\{\s*"
  matches = re.findall(pattern, input_string, re.DOTALL)
  if matches:
    tmp_str = re.sub(pattern, f"}}{spec_line_split}{{", input_string)
    tmp_str_list = tmp_str.split(spec_line_split)
    rtn_list = []
    for str in tmp_str_list:
      try:
        json_data = json.loads(str)
        rtn_list.append(json_data)
      except Exception as e:
        rtn_list.append(str)
    return rtn_list
  else:
    return input_string

def wrap_result(data, raw_ai_result=None):
  """
  封装API返回结果
  :param data: 原始数据
  :return: 封装后的结果
  """
  return {
    'code': 200,
    'data': data,
    'str_data': str(data),
    'raw_ai_result': raw_ai_result
  }

class AiProxy:
  def __init__(self, ai_client, ai_model):
    self.ai_client = ai_client
    self.ai_model = ai_model

  @time_count
  def get_ai_response(self, input_text, sys_role_content=None, is_try_again=False):
    sys_role_content = sys_role_content if sys_role_content else "你是一个善于从文本中提取关键信息的人,请使用中文进行交流"
    completion = self.ai_client.chat.completions.create(
      model=self.ai_model,
      messages=[
        {"role": "system", "content": sys_role_content},
        {"role": "user", "content": input_text}
      ],
      temperature=0.1,
      stream=False,
      max_tokens=4096
    )
    raw_ai_result = completion.choices[0].message.content
    json_data = parse_chat_response(raw_ai_result)
    if isinstance(json_data, str) and not is_try_again:
      json_data = try_parse_json_datas(json_data)
    return wrap_result(json_data, raw_ai_result=raw_ai_result)
